// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <memory>
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/memory/memory_pool.h>
#include <pollux/exec/driver.h>
#include <pollux/exec/memory_reclaimer.h>
#include <pollux/exec/task.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/common/file/temp_directory_path.h>

namespace kumo::pollux::exec::test {

constexpr int64_t KB = 1024L;
constexpr int64_t MB = 1024L * KB;

constexpr uint64_t kMemoryCapacity = 512 * MB;
constexpr uint64_t kMemoryPoolInitCapacity = 16 * MB;

class FakeMemoryReclaimer : public exec::MemoryReclaimer {
 public:
  FakeMemoryReclaimer() : exec::MemoryReclaimer(0) {}

  static std::unique_ptr<MemoryReclaimer> create() {
    return std::make_unique<FakeMemoryReclaimer>();
  }

  void enterArbitration() override {
    auto* driverThreadCtx = driverThreadContext();
    if (driverThreadCtx == nullptr) {
      return;
    }
    auto* driver = driverThreadCtx->driverCtx()->driver;
    ASSERT_TRUE(driver != nullptr);
    if (driver->task()->enterSuspended(driver->state()) != StopReason::kNone) {
      POLLUX_FAIL("Terminate detected when entering suspension");
    }
  }

  void leaveArbitration() noexcept override {
    auto* driverThreadCtx = driverThreadContext();
    if (driverThreadCtx == nullptr) {
      return;
    }
    auto* driver = driverThreadCtx->driverCtx()->driver;
    ASSERT_TRUE(driver != nullptr);
    driver->task()->leaveSuspended(driver->state());
  }
};

struct TestAllocation {
  memory::MemoryPool* pool{nullptr};
  void* buffer{nullptr};
  size_t size{0};

  size_t free() {
    const size_t freedBytes = size;
    if (pool == nullptr) {
      POLLUX_CHECK_EQ(freedBytes, 0);
      return freedBytes;
    }
    POLLUX_CHECK_GT(freedBytes, 0);
    pool->free(buffer, freedBytes);
    pool = nullptr;
    buffer = nullptr;
    size = 0;
    return freedBytes;
  }
};

/// Begins and ends a section where a thread is running but not counted in its
/// Task. Using this, a Driver thread can for example stop its own Task. For
/// arbitrating memory overbooking, the contending threads go suspended and each
/// in turn enters a global critical section. When running the arbitration
/// strategy, a thread can stop and restart Tasks, including its own. When a
/// Task is stopped, its drivers are blocked or suspended and the strategy
/// thread can alter the Task's memory including spilling or killing the whole
/// Task. Other threads waiting to run the arbitration, are in a suspended state
/// which also means that they are instantaneously killable or spillable.
class TestSuspendedSection {
 public:
  explicit TestSuspendedSection(Driver* driver);
  ~TestSuspendedSection();

 private:
  Driver* driver_;
};

std::shared_ptr<core::QueryCtx> newQueryCtx(
    kumo::pollux::memory::MemoryManager* memoryManager,
    melon::Executor* executor,
    int64_t memoryCapacity = kumo::pollux::memory::kMaxMemory,
    const std::string& queryId = "");

std::unique_ptr<memory::MemoryManager> createMemoryManager(
    int64_t arbitratorCapacity = kMemoryCapacity,
    uint64_t memoryPoolInitCapacity = kMemoryPoolInitCapacity,
    uint64_t maxReclaimWaitMs = 5 * 60 * 1'000,
    uint64_t fastExponentialGrowthCapacityLimit = 0,
    double slowCapacityGrowPct = 0);

// Contains the query result.
struct QueryTestResult {
  std::shared_ptr<Task> task;
  RowVectorPtr data;
  core::PlanNodeId planNodeId;
};

core::PlanNodePtr hashJoinPlan(
    const std::vector<RowVectorPtr>& vectors,
    core::PlanNodeId& joinNodeId);

QueryTestResult runHashJoinTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    bool enableSpilling,
    const RowVectorPtr& expectedResult = nullptr);

core::PlanNodePtr aggregationPlan(
    const std::vector<RowVectorPtr>& vectors,
    core::PlanNodeId& aggregateNodeId);

QueryTestResult runAggregateTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    bool enableSpilling,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    const RowVectorPtr& expectedResult = nullptr);

core::PlanNodePtr orderByPlan(
    const std::vector<RowVectorPtr>& vectors,
    core::PlanNodeId& orderNodeId);

QueryTestResult runOrderByTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    bool enableSpilling,
    const RowVectorPtr& expectedResult = nullptr);

core::PlanNodePtr rowNumberPlan(
    const std::vector<RowVectorPtr>& vectors,
    core::PlanNodeId& rowNumberNodeId);

QueryTestResult runRowNumberTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    bool enableSpilling,
    const RowVectorPtr& expectedResult = nullptr);

core::PlanNodePtr topNPlan(
    const std::vector<RowVectorPtr>& vectors,
    core::PlanNodeId& topNodeId);

QueryTestResult runTopNTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    bool enableSpilling,
    const RowVectorPtr& expectedResult = nullptr);

core::PlanNodePtr writePlan(
    const std::vector<RowVectorPtr>& vectors,
    const std::string& outputDirPath,
    core::PlanNodeId& writeNodeId);

QueryTestResult runWriteTask(
    const std::vector<RowVectorPtr>& vectors,
    const std::shared_ptr<core::QueryCtx>& queryCtx,
    bool serialExecution,
    uint32_t numDrivers,
    memory::MemoryPool* pool,
    const std::string& kHiveConnectorId,
    bool enableSpilling,
    const RowVectorPtr& expectedResult = nullptr);
} // namespace kumo::pollux::exec::test
